"""
Agent 核心类

这是整个 Agent 系统的核心，负责协调各个组件的工作。
它展示了一个完整的 AI Agent 应该如何：
1. 接收和理解用户输入
2. 进行推理和决策
3. 调用工具执行任务
4. 管理记忆和上下文
5. 规划和执行复杂任务
"""

import os
import json
import asyncio
from typing import Dict, List, Any, Optional, Union
from dataclasses import dataclass
from enum import Enum

try:
    import openai
except ImportError:
    openai = None

try:
    from rich.console import Console
    from rich.panel import Panel
    from rich.syntax import Syntax
    from rich.markdown import Markdown
except ImportError:
    # 如果没有 rich，使用简单的替代
    class Console:
        def print(self, *args, **kwargs):
            print(*args)
    class Panel:
        def __init__(self, content, **kwargs):
            self.content = content
        def __str__(self):
            return str(self.content)
    class Markdown:
        def __init__(self, content):
            self.content = content
        def __str__(self):
            return str(self.content)

from .reasoning import ReasoningEngine, ReasoningStep, ActionType
from .memory import MemorySystem, MemoryEntry
from .planner import TaskPlanner, Task
from .ollama_client import OllamaOpenAICompatible
from tools.base import Tool, ToolRegistry


class AgentState(Enum):
    """Agent 状态枚举"""
    IDLE = "idle"
    THINKING = "thinking"
    ACTING = "acting"
    PLANNING = "planning"
    ERROR = "error"


@dataclass
class AgentConfig:
    """Agent 配置类"""
    name: str = "DemoAgent"
    provider: str = "openai"  # 支持 'openai' 或 'ollama'
    model: str = "gpt-4"
    temperature: float = 0.7
    max_tokens: int = 2000
    max_reasoning_steps: int = 10
    enable_memory: bool = True
    enable_planning: bool = True
    verbose: bool = True
    
    # Ollama 特定配置
    ollama_base_url: str = "http://localhost:11434"


class Agent:
    """
    AI Agent 核心类
    
    这个类展示了一个完整的 AI Agent 系统的核心架构：
    - 推理引擎：负责分析问题和制定策略
    - 工具系统：提供各种能力和功能
    - 记忆系统：维护上下文和历史记录
    - 任务规划：处理复杂的多步骤任务
    """
    
    def __init__(self, config: AgentConfig):
        self.config = config
        self.state = AgentState.IDLE
        self.console = Console()
        
        # 初始化各个系统组件
        self.reasoning_engine = ReasoningEngine()
        self.memory_system = MemorySystem() if config.enable_memory else None
        self.task_planner = TaskPlanner() if config.enable_planning else None
        self.tool_registry = ToolRegistry()
        
        # 自动注册内置工具
        self._register_built_in_tools()
        
        # 初始化 AI 模型客户端
        self._init_ai_client()
        
        # 对话历史
        self.conversation_history: List[Dict[str, Any]] = []
        
        if config.verbose:
            self._display_welcome()
    
    def _init_ai_client(self):
        """初始化 AI 客户端"""
        # 根据配置选择使用 OpenAI 或 Ollama
        provider = getattr(self.config, 'provider', 'openai')
        
        if provider == 'ollama':
            # 使用 Ollama 本地模型
            ollama_base_url = getattr(self.config, 'ollama_base_url', 'http://localhost:11434')
            try:
                self.client = OllamaOpenAICompatible(
                    base_url=ollama_base_url,
                    model=self.config.model
                )
                if self.config.verbose:
                    print(f"✅ 使用 Ollama 本地模型: {self.config.model}")
            except Exception as e:
                print(f"警告: 初始化 Ollama 客户端失败: {e}，将使用模拟响应")
                self.client = None
        else:
            # 使用 OpenAI API
            api_key = os.getenv("OPENAI_API_KEY")
            if not api_key:
                print("警告: 未设置 OPENAI_API_KEY，将使用模拟响应")
                self.client = None
                return
            
            if openai is None:
                print("警告: 未安装 openai 库，将使用模拟响应")
                self.client = None
                return
                
            try:
                self.client = openai.OpenAI(api_key=api_key)
                if self.config.verbose:
                    print(f"✅ 使用 OpenAI 模型: {self.config.model}")
            except Exception as e:
                print(f"警告: 初始化 OpenAI 客户端失败: {e}，将使用模拟响应")
                self.client = None
    
    def _register_built_in_tools(self):
        """注册内置工具"""
        try:
            from tools.calculator import CalculatorTool
            from tools.search import SearchTool
            from tools.filesystem import FileSystemTool
            from tools.weather import WeatherTool
            
            # 注册工具
            self.tool_registry.register(CalculatorTool())
            self.tool_registry.register(SearchTool())
            self.tool_registry.register(FileSystemTool())
            self.tool_registry.register(WeatherTool())
            
            if self.config.verbose:
                print(f"✅ 已注册 {len(self.tool_registry.list_tools())} 个工具")
                
        except Exception as e:
            if self.config.verbose:
                print(f"⚠️ 注册工具时发生错误: {e}")
    
    def _display_welcome(self):
        """显示欢迎信息"""
        welcome_text = f"""
# 🤖 {self.config.name} Agent 演示系统

这是一个用于教学和演示的 AI Agent 系统。

## 当前配置
- **模型**: {self.config.model}
- **温度**: {self.config.temperature}
- **记忆系统**: {'启用' if self.config.enable_memory else '禁用'}
- **任务规划**: {'启用' if self.config.enable_planning else '禁用'}

## 可用工具
{self._format_available_tools()}

开始与 Agent 对话吧！输入 'help' 查看更多命令。
        """
        
        self.console.print(Panel(
            Markdown(welcome_text),
            title="AI Agent 演示系统",
            border_style="cyan"
        ))
    
    def _format_available_tools(self) -> str:
        """格式化可用工具列表"""
        tools = self.tool_registry.list_tools()
        if not tools:
            return "- 暂无可用工具"
        
        return "\n".join([f"- **{tool}**: {self.tool_registry.get_tool(tool).description}" 
                         for tool in tools])
    
    async def process_input(self, user_input: str) -> str:
        """
        处理用户输入的主要方法
        
        这个方法展示了完整的 Agent 工作流程：
        1. 理解用户意图
        2. 进行推理分析
        3. 制定执行计划
        4. 调用工具执行
        5. 生成响应
        """
        try:
            self.state = AgentState.THINKING
            
            if self.config.verbose:
                self.console.print(f"\n[bold blue]用户:[/bold blue] {user_input}")
                self.console.print("[yellow]🤔 Agent 正在思考...[/yellow]")
            
            # 1. 添加到对话历史
            self.conversation_history.append({
                "role": "user",
                "content": user_input,
                "timestamp": self._get_timestamp()
            })
            
            # 2. 检查是否是特殊命令
            if user_input.lower() in ['help', '帮助']:
                return self._show_help()
            elif user_input.lower() in ['memory', '记忆']:
                return self._show_memory()
            elif user_input.lower() in ['clear', '清空']:
                return self._clear_conversation()
            
            # 3. 构建系统提示
            system_prompt = self._build_system_prompt()
            
            # 4. 获取相关记忆
            relevant_memories = []
            if self.memory_system:
                relevant_memories = await self.memory_system.retrieve_relevant(user_input)
            
            # 5. 进行推理分析
            reasoning_result = await self.reasoning_engine.analyze(
                user_input, 
                self.conversation_history,
                relevant_memories,
                list(self.tool_registry.list_tools())
            )
            
            if self.config.verbose:
                self._display_reasoning(reasoning_result)
            
            # 6. 根据推理结果决定行动
            response = await self._execute_reasoning_plan(reasoning_result)
            
            # 7. 保存到记忆系统
            if self.memory_system:
                await self.memory_system.store(MemoryEntry(
                    content=f"用户问题: {user_input}\nAgent回答: {response}",
                    metadata={"type": "conversation", "reasoning_steps": len(reasoning_result.steps)}
                ))
            
            # 8. 添加到对话历史
            self.conversation_history.append({
                "role": "assistant", 
                "content": response,
                "timestamp": self._get_timestamp(),
                "reasoning": reasoning_result.steps
            })
            
            self.state = AgentState.IDLE
            return response
            
        except Exception as e:
            self.state = AgentState.ERROR
            error_msg = f"处理过程中发生错误: {str(e)}"
            if self.config.verbose:
                self.console.print(f"[red]❌ {error_msg}[/red]")
            return error_msg
    
    def _build_system_prompt(self) -> str:
        """构建系统提示"""
        return f"""你是 {self.config.name}，一个智能的 AI Agent。

你的核心能力：
1. 理解和分析用户需求
2. 使用可用的工具来解决问题
3. 进行逻辑推理和规划
4. 学习和记忆重要信息

可用工具：
{self._format_tools_for_prompt()}

请根据用户的问题进行分析，如果需要使用工具，请明确说明使用什么工具以及原因。
保持回答简洁明了，同时提供有用的信息。
"""
    
    def _format_tools_for_prompt(self) -> str:
        """为提示格式化工具信息"""
        tools = []
        for tool_name in self.tool_registry.list_tools():
            tool = self.tool_registry.get_tool(tool_name)
            tools.append(f"- {tool_name}: {tool.description}")
        return "\n".join(tools) if tools else "暂无可用工具"
    
    async def _execute_reasoning_plan(self, reasoning_result) -> str:
        """执行推理计划"""
        self.state = AgentState.ACTING
        
        # 如果推理结果包含工具调用
        tool_results = []
        for step in reasoning_result.steps:
            if step.action_type == ActionType.TOOL_CALL and step.action_data:
                tool_name = step.action_data.get("tool")
                tool_args = step.action_data.get("args", {})
                
                if self.tool_registry.has_tool(tool_name):
                    if self.config.verbose:
                        self.console.print(f"[green]🔧 调用工具: {tool_name}[/green]")
                    
                    tool = self.tool_registry.get_tool(tool_name)
                    result = await tool.execute(**tool_args)
                    tool_results.append({
                        "tool": tool_name,
                        "args": tool_args,
                        "result": result
                    })
        
        # 生成最终响应
        return await self._generate_response(reasoning_result, tool_results)
    
    async def _generate_response(self, reasoning_result, tool_results: List[Dict]) -> str:
        """生成最终响应"""
        # 如果没有 AI 客户端，使用模拟响应
        if self.client is None:
            return self._generate_mock_response(reasoning_result, tool_results)
        
        # 构建包含推理过程和工具结果的提示
        messages = [
            {"role": "system", "content": self._build_system_prompt()},
        ]
        
        # 添加最近的对话历史
        recent_history = self.conversation_history[-5:]  # 最近5轮对话
        messages.extend([
            {"role": msg["role"], "content": msg["content"]} 
            for msg in recent_history
        ])
        
        # 添加推理过程和工具结果
        reasoning_summary = "\n".join([
            f"步骤 {i+1}: {step.thought}" 
            for i, step in enumerate(reasoning_result.steps)
        ])
        
        tool_summary = ""
        if tool_results:
            tool_summary = "\n工具执行结果:\n" + "\n".join([
                f"- {result['tool']}: {result['result']}"
                for result in tool_results
            ])
        
        final_prompt = f"""
基于以下推理过程和工具执行结果，请生成一个完整、有用的回答：

推理过程：
{reasoning_summary}

{tool_summary}

请提供一个自然、友好的回答，整合所有相关信息。
"""
        
        messages.append({"role": "user", "content": final_prompt})
        
        try:
            # 检查是否为 Ollama 客户端（异步调用）
            if hasattr(self.client, 'chat') and hasattr(self.client.chat, 'create'):
                # OpenAI 或 Ollama 的异步调用
                if asyncio.iscoroutinefunction(self.client.chat.create):
                    # Ollama 异步调用
                    response = await self.client.chat.create(
                        model=self.config.model,
                        messages=messages,
                        temperature=self.config.temperature,
                        max_tokens=self.config.max_tokens
                    )
                else:
                    # OpenAI 同步调用
                    response = self.client.chat.completions.create(
                        model=self.config.model,
                        messages=messages,
                        temperature=self.config.temperature,
                        max_tokens=self.config.max_tokens
                    )
            else:
                # 备用方案
                response = await self.client.chat.create(
                    model=self.config.model,
                    messages=messages,
                    temperature=self.config.temperature,
                    max_tokens=self.config.max_tokens
                )
            
            return response.choices[0].message.content.strip()
            
        except Exception as e:
            return f"生成响应时发生错误: {str(e)}"
    
    def _generate_mock_response(self, reasoning_result, tool_results: List[Dict]) -> str:
        """生成模拟响应（当没有 AI 客户端时）"""
        response_parts = []
        
        # 基于推理结果生成响应
        if reasoning_result.steps:
            first_step = reasoning_result.steps[0]
            if "计算" in first_step.thought:
                response_parts.append("我理解您需要进行数学计算。")
            elif "搜索" in first_step.thought:
                response_parts.append("我将为您搜索相关信息。")
            elif "文件" in first_step.thought:
                response_parts.append("我将执行文件操作。")
            else:
                response_parts.append("我已经分析了您的请求。")
        
        # 添加工具结果
        if tool_results:
            response_parts.append("\n执行结果：")
            for result in tool_results:
                response_parts.append(f"- {result['tool']}: {result['result']}")
        
        # 默认响应
        if not response_parts:
            response_parts = ["我已经理解了您的请求，这是一个模拟响应。"]
        
        response_parts.append("\n（注意：这是模拟响应，请配置 API 密钥以获得真实 AI 响应）")
        
        return "".join(response_parts)
    
    def _display_reasoning(self, reasoning_result):
        """显示推理过程"""
        if not reasoning_result.steps:
            return
        
        reasoning_text = "## 🧠 推理过程\n\n"
        for i, step in enumerate(reasoning_result.steps, 1):
            reasoning_text += f"**步骤 {i}**: {step.thought}\n"
            if step.action_type:
                reasoning_text += f"- 行动类型: {step.action_type}\n"
            if step.confidence_score:
                reasoning_text += f"- 信心指数: {step.confidence_score:.1f}/10\n"
            reasoning_text += "\n"
        
        self.console.print(Panel(
            Markdown(reasoning_text),
            title="推理分析",
            border_style="yellow"
        ))
    
    def _show_help(self) -> str:
        """显示帮助信息"""
        help_text = """
## 🎯 可用命令

- **help** / **帮助**: 显示此帮助信息
- **memory** / **记忆**: 查看记忆系统状态
- **clear** / **清空**: 清空对话历史

## 🛠️ 可用工具

""" + self._format_available_tools() + """

## 💡 使用提示

1. 直接输入问题或请求，Agent 会自动分析并选择合适的工具
2. 对于复杂任务，Agent 会自动进行任务分解和规划
3. Agent 会记住对话历史，可以进行上下文相关的对话
4. 输入 'memory' 可以查看 Agent 学到了什么

开始提问吧！
"""
        return help_text
    
    def _show_memory(self) -> str:
        """显示记忆状态"""
        if not self.memory_system:
            return "记忆系统未启用"
        
        entries = self.memory_system.get_recent_entries(5)
        if not entries:
            return "暂无记忆内容"
        
        memory_text = "## 🧠 最近的记忆\n\n"
        for i, entry in enumerate(entries, 1):
            memory_text += f"**记忆 {i}**:\n{entry.content[:100]}...\n\n"
        
        return memory_text
    
    def _clear_conversation(self) -> str:
        """清空对话历史"""
        self.conversation_history = []
        return "✅ 对话历史已清空"
    
    def _get_timestamp(self) -> str:
        """获取时间戳"""
        import datetime
        return datetime.datetime.now().isoformat()
    
    def register_tool(self, tool: Tool):
        """注册新工具"""
        self.tool_registry.register(tool)
        if self.config.verbose:
            self.console.print(f"[green]✅ 工具已注册: {tool.name}[/green]")
    
    def get_status(self) -> Dict[str, Any]:
        """获取 Agent 状态"""
        return {
            "state": self.state.value,
            "config": self.config.__dict__,
            "conversation_length": len(self.conversation_history),
            "available_tools": list(self.tool_registry.list_tools()),
            "memory_entries": len(self.memory_system.entries) if self.memory_system else 0
        }